#include <iostream>
#include <vector>
#include <ctime>

#include "tensor.h"
#include "covariance.h"
#include "mtrand.h"
#include "dataset.h"
#include "pdf.h"
#include <vector>

#include "TFile.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TF1.h"
#include "TF2.h"
#include "TGraphErrors.h"
#include "TGraph2DErrors.h"
#include "make_list.h"

using namespace std;
typedef vector<double> V1;

//initialize random number generator
MTRand_open irand(time(0));

const double p0=1.25658381326133917e+00;
const double p1=2.97897717730167066e+01;
const double p2=2.68831342433453706e+01;
const double p3=-1.76414975998154655e+00;
const double norm=1.002707104*pow(2*3.1415926, -0.5);

//function to test PDF estimation
inline double testfunc(double x)
{
  return p0*x*(1-x)*exp(-p1*x*x*x+p2*x*x+p3*x);
}

//2D function
inline double testfunc2D(double x, double y)
{
  //x goes from 0 to 2
  //y goes from 0 to 1
  double tx=(2-x);
  tx*=tx;
  tx*=tx;

  double ty=(1-y);
  ty*=ty;

  return 0.25*testfunc(0.5*x)*testfunc(y)+ 45*x*tx*y*ty/16.;
}



//take a 2D function and generate MC
pair<V1,V1> MC2D( double (*func)(double, double), int nevt=1000)
{
  //sample the distribution first
  double abs_max=2.4;

  abs_max*=1.10;
  
  V1 resultx;
  V1 resulty;
  for(int i=0; i<nevt; ++i)
    {
      bool success=false;
      while(!success)
	{
	  double numx=irand()*2.0;
	  double numy=irand();
	  double num_reject=irand()*abs_max;
	  if(func(numx, numy) > num_reject)
	    {
	      resultx.push_back(numx);
	      resulty.push_back(numy);
	      success=true;
	      break;
	    }
	}
    }
  return pair<V1, V1>(resultx, resulty);
}



//take a function and generate MC
V1 MC( double (*func)(double), int nevt=1000)
{
  //sample the distribution first
  V1 samples;
  V1 samples_max;
  for(double i=0; i<1.0; i+=0.05)
    {samples.push_back(func(i));}

  for(int i=1; i<samples.size()-1; ++i)
    {
      samples_max.push_back
	(samples[i]+0.05*max(samples[i]-samples[i-1], 
			     samples[i]-samples[i+1]));

    }

  double abs_max=*(max_element(samples_max.begin(), 
			      samples_max.end()));

  V1 result;
  for(int i=0; i<nevt; ++i)
    {
      bool success=false;
      while(!success)
	{
	  double num=irand();
	  double num_reject=irand()*abs_max;
	  if(func(num) > num_reject)
	    {
	      result.push_back(num);
	      success=true;
	      break;
	    }
	}
    }
  return result;
}



int main()
{
  TF1 f("f", "[0]*x*(1-x)*exp(-[1]*x*x*x+[2]*x*x+[3]*x)", -0.1, 1.1);
  f.SetParameters(p0, p1, p2, p3);

  TH1D MCf("MCf", "MCf", 100, 0, 1.0);
  vector<double> x,y,yerr;
  
  //first initialize DataSet
  //with number of variables, min and max
  double xmin[]={0};
  double xmax[]={1.0};

  DataSet mydata(1, xmin, xmax);

  //generate some MC to fill mydata
  V1 result= MC(testfunc, 2000);

  for(int i=0; i<result.size(); i++)
    {
      mydata.Fill(make_list<double>(result[i]));
      MCf.Fill(result[i]);
    }


  //now generate PDF
  PDF mypdf=mydata.ComputePDF();

  for(double i=0.01; i<1.0; i+=0.07)
    {
      double err=0;
      x.push_back(i);
      double idouble[]={i};
      valarray<double> iary(idouble,1);
      y.push_back(mypdf(iary, &err));
      yerr.push_back(err);
    }

  TGraphErrors g(x.size(), &x[0], &y[0], NULL, &yerr[0]);
  g.SetName("Kernel");
  g.SetTitle("Kernel");


  //now test 2D capabilities
  //let us define a 2D function in root
  TF2 f2("f2", "0.25*([0]*y*(1-y)*exp(-[1]*y*y*y+[2]*y*y+[3]*y))*([0]*0.5*x*(1-0.5*x)*exp(-[1]*0.125*x*x*x+[2]*0.25*x*x+[3]*0.5*x))+45*x*(2-x)^4*y*(1-y)^2/16", 0, 2.0, 0, 1.0);
  
  f2.SetParameters(p0, p1, p2, p3);
  TH2D MCf2D("MCf2D", "MCf2D", 20, 0, 2.0, 20, 0, 1.0);

  double xmin2D[]={0, 0};
  double xmax2D[]={2.0, 1.0};
  double precision[]={0.005, 0.005};
  DataSet mydata2D(2, xmin2D, xmax2D, precision);

  //vector for graph2D
  vector<Double_t> x2D,y2D,z,zerr;

  cout<<"running 2D MC"<<endl;
  pair<V1,V1> result2D= MC2D(testfunc2D, 10000000);
  cout<<"filling 2D dataset"<<endl;
  for(int i=0; i<result2D.first.size(); ++i)
    {
      MCf2D.Fill(result2D.first[i], result2D.second[i]);
      mydata2D.Fill(make_list<double>(result2D.first[i], result2D.second[i]));
    }
  
  //now generate PDF
  cout<<"Computing PDF"<<endl;
  PDF mypdf2D=mydata2D.ComputePDF();

  cout<<"Done Computing PDF, evaluating"<<endl;

  TH1D num_sig("sig", "sig", 100, -5, 5);

  for(double i=0; i<2.0; i+=0.08)
    for(double j=0; j<1.0; j+=0.08)
    {
      double err=0;
      x2D.push_back(i);
      y2D.push_back(j);
      z.push_back(mypdf2D(make_list<double>(i,j), &err));

      double sig=0;
      if(err>0)
	sig=(z.back()-testfunc2D(i,j))/err;
      else
	{
	  double diff=(z.back()-testfunc2D(i,j));
	  if(diff>0)
	    sig=5;
	  else if(diff<0)
	    sig=-5;
	  else
	    sig=0;
	}

      /*
      cout<<"(actual, estimate, error, sig): ("
	  <<testfunc2D(i,j)<<","<<z.back()<<","<<err<<","<<sig<<")"<<endl;
      */
      num_sig.Fill(sig);
      zerr.push_back(err);
    }


      cout<<"Done evaluating"<<endl;
  

  TGraph2DErrors g2D(x2D.size(), &x2D[0], &y2D[0], &z[0], 0, 0, &zerr[0]);
  g2D.SetName("g2D");
  g2D.SetTitle("g2D");



  //take y slices
  vector<double> xslice;
  vector<double> yslice_actual;  
  vector<double> yslice;  
  vector<double> yslice_err;  
  
  for(double i=0.01; i<1.0; i+=0.03)
    {
      double err;
      xslice.push_back(i);
      yslice_actual.push_back(testfunc2D(0.67, i));
      yslice.push_back(mypdf2D(make_list<double>(0.67,i), &err));
      yslice_err.push_back(err);
    }
  
  TGraphErrors gslice(xslice.size(), &xslice[0], &yslice[0], 0, &yslice_err[0]);
  TGraph gslice_actual(xslice.size(), &xslice[0], &yslice_actual[0]);
  
  gslice.SetName("pdf_slice");
  gslice_actual.SetName("actual_slice");

  //write to root file
  TFile outfile("outfile.root", "RECREATE");
  f.Write();
  g.Write();
  MCf.Scale(100/MCf.Integral());
  MCf.Write();
  
  num_sig.Write();
  f2.Write();
  MCf2D.Write();
  g2D.Write();

  gslice.Write();
  gslice_actual.Write();

  outfile.Close();

}
